import networkx as nx
import random
import math
import matplotlib.pyplot as plt

from networkx import Graph
from networkx.generators.community import stochastic_block_model
import numpy as np

class SA:
    def __init__(self) -> None:
        self.G = nx.Graph()
    
    def run(self):
        self.randomEdgeWithNodeNum()

        initial_temp = 1.0
        cooling_rate = 0.99
        best_partition, best_modularity = self.simulated_annealing(initial_temp, cooling_rate)

        # print("Best Partition:", best_partition)
        # print("Best Modularity:", best_modularity)
        # print(self.mapToListIntInt(best_partition))
        self.buildOverlappedNetwork(best_partition)

        # self.showCommunity(self.mapToListIntInt(best_partition))

    def buildOverlappedNetwork(self, partition: dict[int, int]) -> tuple[Graph, dict[int, int]]:
        # 根据 overlapModularityQ 的准则，构建一个具有交叉节点的网络
        # 1. 遍历一个簇的相邻节点，计算将每个相邻节点各自加入到这个簇的模块度增加值
        # 2. 取Top N模块度的节点加入到这个簇，使簇的节点数量为15-20
        cluster_list = self.mapToListIntInt(partition)
        cluster_list = [cluster for cluster in cluster_list if len(cluster) != 0]
        print(cluster_list)
        for cluster in cluster_list:
            # 计算每个节点的模块度增加值
            modularity_increase = {}
            for node in cluster:
                # 获取节点的相邻节点
                neighbors = list(self.G.neighbors(node))
                for neighbor in neighbors:
                    if neighbor not in cluster:
                        # 将邻居节点加入到簇中，计算模块度增加值
                        old_cluster_list = cluster_list.copy()
                        cluster.append(neighbor)
                        modularity_increase[neighbor] = self.overlapModularityQ(cluster_list) - self.overlapModularityQ(old_cluster_list)
                        cluster.remove(neighbor)
            # 将模块度增加值排序，取Top N
            sorted_modularity_increase = sorted(modularity_increase.items(), key=lambda x: x[1], reverse=True)
            print("modularity_increase: ", modularity_increase)
            for node, _ in sorted_modularity_increase:
                if node not in cluster:
                    cluster.append(node)
            
            print("cluster: ", cluster)
            # 截取前20个节点
            if len(cluster) > 20:
                cluster = cluster[:20]

        print("cluster_list: ", cluster_list)
        # print(cluster_list)

    # 传入图和社区划分结果，计算重叠模块度
    def overlapModularityQ(self, communities: list[list[int]]) -> float:
        # start = time.time()
        # 网络边数量
        edge_num = self.G.number_of_edges()
        def delta_cv(cluster: list[int], node: int):
            # 计算节点 v 归属于簇 c 的概率
            degree_cv = sum([1 for n in cluster if self.G.has_edge(node, n)]) # 计算节点node有多少条边与cluster中节点相连
            return 1 / (1 + math.exp(-120 * degree_cv + 60))
        def l_uv(u: int, v: int):
            # 计算节点 u 和 v 是否相邻
            return 1 if self.G.has_edge(u, v) else 0
        def k_u(u: int):
            # 计算节点 u 的度数
            return self.G.degree(u)
        modularity = 0
        for cluster in communities:
            for u in self.G.nodes:
                for v in self.G.nodes:
                    tmp_cv = delta_cv(cluster, u) * delta_cv(cluster, v) * (l_uv(u, v) - k_u(u) * k_u(v) / (2 * edge_num))
                    modularity += 1 / (2 * edge_num) * tmp_cv
        return modularity

    def showCommunity(self, communities: list[list[int]]):
        community_map = {}
        for idx, community in enumerate(communities):
            for node in community:
                community_map[node] = idx
        color_map = [community_map[node] for node in self.G.nodes()]
        pos = {}
        radius = 1
        angles = np.linspace(0, 2 * np.pi, len(communities), endpoint=False)
        community_centers = radius * np.column_stack([np.cos(angles), np.sin(angles)])
        for idx, community in enumerate(communities):
            subgraph = self.G.subgraph(community)
            community_pos = nx.spring_layout(subgraph, scale=0.5)  # scale小一些以保持社区紧凑
            community_pos = {node: pos + community_centers[idx] for node, pos in community_pos.items()}
            pos.update(community_pos)
        plt.figure(figsize=(8, 8))
        nx.draw(self.G, pos, node_color=color_map, with_labels=True, node_size=100, cmap=plt.cm.jet)
        nx.draw_networkx_edges(self.G, pos, alpha=0.3)
        plt.axis('off')  # 隐藏坐标轴
        plt.show()

    def mapToListIntInt(self, map: dict[int, int]) -> list[list[int]]:
        result = {}
        for key, value in map.items():
            if value not in result:
                result[value] = []
            result[value].append(key)
        return list(result.values())

    def getModularityGap(self, path_length: int) -> float:
        if(path_length < 20):
            return 0.08
        if(path_length < 50):
            return 0.06
        if(path_length < 80):
            return 0.04
        if(path_length < 100):
            return 0.02
        if(path_length < 120):
            return -0.02
        if(path_length < 150):
            return -0.05
        if(path_length < 200):
            return -0.08
        return -0.1

    def getClusterNodeNum(self, partition: dict, nodeID: int) -> int:
        result = 0
        for _, value in partition.items():
            if value == nodeID:
                result += 1
        return result

    def simulated_annealing(self, initial_temp, cooling_rate, iteration_num=1000):
        # 将每个节点赋予一个社区
        partition = {node: i for i, node in enumerate(self.G.nodes)}
        # 计算初始划分的模块度
        modularity = nx.algorithms.community.modularity(self.G, [set(self.G.nodes)])
        best_partition = partition
        best_modularity = modularity
        temp = initial_temp
        for _ in range(iteration_num):
            for node in self.G.nodes:
                # 获取当前已有的社区列表
                community_choices = list(set(partition.values()))
                # 计算每个社区的节点数
                community_sizes = {community: list(partition.values()).count(community) for community in community_choices}
                # 随机选择一个新的社区，但是不要选择那些节点数超过16个节点的社区
                # 从社区列表中移除节点数超过16的社区
                community_choices = [community for community in community_choices if community_sizes[community] <= 15]
                # 随机选择一个新的社区
                new_community = random.choice(community_choices)
                # 记录旧的社区
                old_community = partition[node]
                # 将节点移动到新的社区
                partition[node] = new_community

                # 计算节点 node 到 new_community 之间的最短路径长度
                target_nodes = [n for n, com in partition.items() if com == new_community and n != node]
                shortest_path_lengths = {target: nx.shortest_path_length(self.G, source=node, target=target, weight='weight') for target in target_nodes}

                # 获取最短的shortest_path_lengths
                shortest_path_length = min(shortest_path_lengths.values()) if len(shortest_path_lengths) != 0 else 0

                # 计算新的划分的模块度
                new_modularity = nx.algorithms.community.modularity(self.G, list(set([frozenset([node for node, community in partition.items() if community == com]) for com in set(partition.values())]))) + self.getModularityGap(shortest_path_length)
                delta_modularity = new_modularity - modularity
                if delta_modularity > 0 or math.exp(delta_modularity / temp) > random.random():
                    modularity = new_modularity
                    if new_modularity > best_modularity:
                        print("New Modularity:", new_modularity)
                        best_partition = partition.copy()
                        best_modularity = new_modularity
                else:
                    partition[node] = old_community  # 撤回变更
            temp *= cooling_rate
        return best_partition, best_modularity

    def randomWeight(self) ->int :
        return random.randint(20, 100)

    def randomEdgeWithNodeNum(self):
        node_num = 40
        edge_num = 120
        self.G.add_nodes_from(range(1, node_num+1))
        while self.G.number_of_edges() < edge_num:
            node1 = random.choice(range(1, node_num+1))
            node2 = random.choice(range(1, node_num+1))
            if node1 != node2 and not self.G.has_edge(node1, node2):
                self.G.add_edge(node1, node2, weight=self.randomWeight())
        for node in self.G.nodes:
            if self.G.degree(node) == 0:
                random_node = random.choice(range(node_num))
                self.G.add_edge(node, random_node, weight=self.randomWeight())

    def addEdgeWithCertain(self):
        self.G.add_weighted_edges_from([(1,3,20),(1,4,43),(2,3,37),(2,4,41),(4,5,29),(4,6,21),(5,7,33),(6,7,23)])

    def addGWithClusterShape(self):
        sizes = [20, 25, 30]  # n个簇的节点数
        p_intra = 0.3  # 簇内的节点之间的连边概率
        p_inter = 0.01  # 簇间的节点之间的连边概率
        p_matrix = [[p_intra if i == j else p_inter for j in range(len(sizes))] for i in range(len(sizes))]
        self.G.add_edges_from(stochastic_block_model(sizes, p_matrix, seed=0).edges())
        # 为G的每条边分配随机的权重
        for u, v in self.G.edges:
            self.G[u][v]['weight'] = self.randomWeight()

sa = SA()
sa.run()
